"""
@Description :   Qwen2.5-VL 模型 API
@Author      :   tqychy 
@Time        :   2025/08/24 16:10:05
"""
import sys

sys.path.append("./nets")
import base64
import json
import os

from api_key import api_key
from openai import OpenAI


class Qwen2_5_VL_API:
    def __init__(self, *args, model_path):
        self.cfg, self.logger = args
        self.model_name = model_path
        self.model = OpenAI(
            api_key=api_key("aliyun"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

    def __call__(self, inputs):
        completion = self.model.chat.completions.create(model=self.model_name, messages=inputs)
        response = json.loads(completion.model_dump_json())
        return response["choices"][0]["message"]["content"]
    
    @staticmethod
    def make_prompt(image_path, sentence, category, mode):
        if mode == "rec":
            prompt = f"""
            <image>
            Here is a description of the objects in the figure. 
            Please enclose the corresponding positions using coordinate boxes. 
            Examples of coordinate value formats: [x1, y1, x2, y2]
            Please only output boxes in a fixed format like:
            ```json\n[\n\t{{"bbox_2d": [x1, y1, x2, y2], "label": "..."}}\n]\n```
            and do not output any unnecessary content.

            Description:\n{sentence}
            """
        elif mode == "detect":
            prompt = f"""
            <image>
            Here is a series of category labels separated by "." and an image. Please identify ALL objects in the image corresponding to the mentioned category labels, and meet the following requirements:
            1. There may be multiple objects for each category label; please identify ALL of them.
            2. Return a JSON-formatted list of dictionaries. Each dictionary must contain the "bbox_2d", "label" and "score" of the identified object. An example of the output format is: [{{"bbox_2d": [456, 249, 507, 324], "label": "Dog", "score": 0.67}}, {{"bbox_2d": [495, 261, 528, 322], "label": "Dog", "score": 0.84}}, {{"bbox_2d": [392, 323, 447, 401], "label": "Monkey", "score": 0.46}}]. Do not output any extra information.
            3. Do not output any extra information.
            4. The returned category labels must EXACTLY match the input category labels.
            Categories: {category}
            """
        else:
            raise ValueError(f"Unknown task type! {mode}")

        with open(image_path, "rb") as img_file:
            img_base = base64.b64encode(img_file.read()).decode("utf-8")
        inputs = [
            {
                "role": "user",
                "content": [
                    {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64, {img_base}"
                        },
                    },
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ]
        return inputs
    
    @staticmethod
    def covert_formatted_bbox(bbox, image_shape):
        """
        将 [x1, y1, x2, y2] 格式的 bbox 转换成 [x_min, y_min, width, height] 格式的​
        """
        x1, y1, x2, y2 = bbox
        x_min = min(x1, x2)
        y_min = min(y1, y2)
        width = max(x1, x2) - x_min
        height = max(y1, y2) - y_min
        return [x_min, y_min, width, height]
    
if __name__ == "__main__":
    image_path = "./dataset/scripts/refcoco/data/images/mscoco/images/train2014/COCO_train2014_000000000009.jpg"
    with open(image_path, "rb") as img_file:
        img_base = base64.b64encode(img_file.read()).decode("utf-8")
    inputs = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64, {img_base}"
                        },
                },
                {
                    "type": "text",
                    "text": "请描述这个图片"
                }
            ]
        }
    ]

    model = Qwen2_5_VL_API(None, None)
    print(model(inputs))